import os
import os.path as osp
import shutil
import random
random.seed(0)
BASEDIR = osp.dirname(osp.abspath(__file__))
train_val_rate = 0.8
print(train_val_rate)

IMGDIR = osp.join(BASEDIR, 'cifar-images/data_batch_5_dir')
imgnames = [name for name in os.listdir(IMGDIR) if name.endswith('.png')]
random.shuffle(imgnames)
classnums = sorted(list(set([name.split('_')[0] for name in imgnames])))
print(imgnames)
print(len(imgnames))
print(classnums)
print(len(classnums))

OUTDIR = osp.join(BASEDIR, 'cifar-trainval/data_batch_5_dir')
OUTtrain = osp.join(OUTDIR, 'train')
OUTval = osp.join(OUTDIR, 'val')
for cn in classnums:
    traindir = osp.join(OUTtrain, cn)
    valdir = osp.join(OUTval, cn)
    for DIR in [traindir, valdir]:
        if not osp.exists(DIR):
            os.makedirs(DIR)

# grouping
clsdic = {}
for name in imgnames:
    clsnum = name.split('_')[0]
    if clsnum not in clsdic:
        clsdic[clsnum] = []
    clsdic[clsnum].append(name)

# copying
for clsnum in clsdic:
    train_max = len(clsdic[clsnum]) * train_val_rate
    print(train_max)
    for i, name in enumerate(clsdic[clsnum]):
        src_path = osp.join(IMGDIR, name)
        if i < train_max:
            dst_path = osp.join(OUTtrain, clsnum, name)
        else:
            dst_path = osp.join(OUTval, clsnum, name)
        shutil.copy(src_path, dst_path)
        print(clsnum, name)
